import scipy
from scipy.stats import linregress
import numpy
from numpy import asarray
from numpy import linspace
import matplotlib
matplotlib.use('agg')
from matplotlib import pyplot as PLT
from CalTorque import CalTorque
import os
import sys
from datetime import datetime
from scipy.interpolate import UnivariateSpline

class Velocity:
    ''' enter explanation'''
    def __init__(self, pos, time, run):
        self.pos = pos
        self.time = time
        self.run = run
        self.IntegratedVelocity()

    def VelocityBeforeS5(self):
        max_d = 400
        bins = 400
        
        x_inputs = numpy.zeros(bins)
        X = numpy.linspace(0,max_d,bins)
        Y = numpy.zeros(bins)
        
        # bin data
        for j,x in enumerate(X):
              pos2 = []
              x0 = x-(max_d/bins)/2.
              x1 = x+(max_d/bins)/2.
              for i, p in enumerate(self.time):
                 if p>x0 and p<=x1:
                    pos2.append(self.pos[i])
              if len(pos2)>0:
                 if max(pos2) == 0: continue
                 Y[j] += max(pos2)
                 x_inputs[j]+=1.
        for i,n in enumerate(x_inputs):
           if n==0: continue
           Y[i] = Y[i]/n
    
        
        # slice data to that surrounding S5
        dt = []
        dp = []
        for i, k in enumerate(self.pos):
            if k < 380 and k > 350:
                dp.append(k)
                dt.append(self.time[i])
        
        # fit the data surrounding S5
        linfit = linregress(dt, dp)
        velAtS5 = linfit[0]
        
        lin = []
        for object in dt:
            lin.append(object*velAtS5 + linfit[1])
        
        '''
        # plot fit
        maxx = max(time)
        minx = max(time)*0.07
        fig = PLT.figure(figsize = (15, 8), dpi = 150)
        box = [0.14, 0.14, 0.76, 0.76]
        ax = fig.add_axes(box)
        ax.set_ylabel('Position (cm)')
        ax.set_xlabel('Time (s)')
        PLT.title('velocity spline test')
        ax.plot(time, pos, '-', color = 'blue')
        ax.plot(dt, lin, '-', color = 'red')
        if pos[-1] > 350:
          ax.hlines(388.0, minx, maxx, color='0.82',linewidth=4)
          ax.annotate('S5',xy=(minx*.7,388.0),color='0.75')
        fig.savefig('vel2/' + run + '_VelTest.png')
        PLT.close()'''
        
        return velAtS5
    
    
    def IntegratedVelocity(self):
        timeChop = []
        posChop = []
        
        for i, p in enumerate(self.pos):
            ## TODO more robust way of doing this
            if p == self.pos[0] or p == self.pos[-1] or p >= 385: continue
            posChop.append(p)
            timeChop.append(self.time[i])
        #print timeChop
        #print posChop
        
        spl = UnivariateSpline(timeChop, posChop)
        ys = spl(timeChop)
        
        #print ys
        #print spl(timeChop)
        
        d = spl.derivative()
        yprime = d(timeChop)
        
        fig = PLT.figure(figsize = (15, 8), dpi = 150)
        box = [0.14, 0.14, 0.76, 0.76]
        ax = fig.add_axes(box)
        ax2 = ax.twinx()
        ax.set_ylabel('Position (cm)', color = 'blue')
        ax.set_xlabel('Time (s)')
        ax2.set_ylabel('Velocity (cm/s)', color = 'green')
        PLT.title('Integrated Velocity Test')
        ax.plot(timeChop, posChop, '-', color = 'blue', label = "Position vs Time data")
        ax.plot(x, ys, '-', color = 'red', label = "Spline")
        ax2.plot(timeChop, yprime, '-', color = 'green', label = "Derivative of Spline - Velocity")
        h, l = ax.get_legend_handles_labels()
        q, t = ax2.get_legend_handles_labels()
        ax.set_ylim([-25,400])
        ax.legend(h, l, 'upper left')
        ax2.legend(q, t, 'upper center')
        fig.savefig('vel/runs/IntegratedVelocityTest2' + self.run + '.png')
        PLT.close()
        
        return yprime

    def AverageVelocity(self):
        velocity = self.IntegratedVelocity()
        #print "-> 1", velocity
        counter = 0
        velocitySum = 0
        for v in velocity:
            counter += 1
            velocitySum += v
        #print velocitySum / counter
        return velocitySum / counter

    def MaximumVelocity(self):
        velocity = self.IntegratedVelocity()
        #print "-> 2", velocity
        for i, v in enumerate(velocity):
            if i == 0:
                maxVelocity = v
            if v > maxVelocity:
                maxVelocity = v
        #print maxVelocity
        return maxVelocity

    def MedianVelocity(self):
        velocity = self.IntegratedVelocity()
        #print "-> 3", velocity
        velocity.sort()
        i = len(velocity)/2.
        if len(velocity) % 2 == 0:
           # print (velocity[i] + velocity[i+1]) /2.
            return (velocity[i] + velocity[i+1]) /2.
        else:
            #print velocity[i]
            return velocity[i]


def main(argv = sys.argv):
    
    #data = CalTorque('data/Live Graph 05-Nov-2012_in.xlsx')
    #print data['position']
    
    #vel = Velocity(data['pos_offsetfix'], data['dt'], data['run'])
    #print vel.AverageVelocity()
    #print vel.MaximumVelocity()
    #print vel.MedianVelocity()
    #print vel.AverageVelocity(data['pos_offsetfix'], data['dt'], data['run'])
    
    
    dp = 'data/'
    ls = os.listdir(dp)
    for i in ls:
        if i.split('.')[-1] not in ['xlsx', 'xls']: continue
        if 'out' in i: continue
        elif 'in' not in i: continue
        print i

        data = CalTorque(dp + i)
        vel = Velocity(data['pos_offsetfix'], data['dt'], data['run'])
        vel.IntegratedVelocity()
    
    

if __name__ == "__main__":
   main(sys.argv)

